import numpy as np
import matplotlib.pyplot as plt

from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.datasets import cifar10
from sklearn.preprocessing import LabelBinarizer
from sklearn.metrics import classification_report


def main(save_loss_data: bool = True):
    print('[info]:加载CIFAR-10数据集中...')

    # 数据加载
    (train_x, train_y), (test_x, test_y) = cifar10.load_data()

    # 数据归一化
    train_x = train_x.astype('float') / 255.0
    test_x = test_x.astype('float') / 255.0

    # 数据重塑
    train_x = train_x.reshape((train_x.shape[0], 3072))
    test_x = test_x.reshape((test_x.shape[0], 3072))

    # 标签向量化
    lb = LabelBinarizer()
    train_y = lb.fit_transform(train_y)
    test_y = lb.transform(test_y)

    label_names = ["airplane", "automobile", "bird", "cat", "deer",
                   "dog", "frog", "horse", "ship", "truck"]
    # 构建网络
    print('[info]:构建网络中..')
    model = Sequential()
    model.add(Dense(1024, input_shape=(3072,), activation='relu'))
    model.add(Dense(512, activation='relu'))
    model.add(Dense(10, activation='softmax'))

    # 训练网络
    print('[info]:开始训练...')
    sgd = SGD(0.01)
    model.compile(loss='categorical_crossentropy',
                  optimizer=sgd,
                  metrics=['accuracy'])
    record = model.fit(train_x, train_y,
                       validation_data=(test_x, test_y),
                       epochs=100, batch_size=32)

    # 评估网络
    print('[info]:评估网络中...')
    predictions = model.predict(test_x, batch_size=32)
    print(classification_report(test_y.argmax(1),
                                predictions.argmax(1),
                                target_names=label_names))

    # 画图
    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, 100), record.history["loss"], label="train_loss")
    plt.plot(np.arange(0, 100), record.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, 100), record.history["accuracy"], label="train_acc")
    plt.plot(np.arange(0, 100), record.history["val_accuracy"], label="val_acc")
    plt.title("Training Loss and Accuracy")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend()
    if save_loss_data:
        plt.savefig("cifar10数据集NN网络损失及准确率变化图")


if __name__ == '__main__':
    main()
